import torch.nn as nn
from .util import init

"""CNN 模块和工具。"""

class Flatten(nn.Module):
    """
    展平层，将输入张量展平为二维张量
    """
    def forward(self, x):
        # 将输入张量展平，保留批次维度
        return x.view(x.size(0), -1)


class CNNLayer(nn.Module):
    """
    CNN层，用于处理图像输入
    """
    def __init__(self, obs_shape, hidden_size, use_orthogonal, use_ReLU, kernel_size=3, stride=1):
        """
        初始化CNN层
        
        参数:
            obs_shape: 观察空间的形状
            hidden_size: 隐藏层大小
            use_orthogonal: 是否使用正交初始化
            use_ReLU: 是否使用ReLU激活函数
            kernel_size: 卷积核大小
            stride: 卷积步长
        """
        super(CNNLayer, self).__init__()

        # 根据参数选择激活函数
        active_func = [nn.Tanh(), nn.ReLU()][use_ReLU]
        # 根据参数选择初始化方法
        init_method = [nn.init.xavier_uniform_, nn.init.orthogonal_][use_orthogonal]
        # 计算增益
        gain = nn.init.calculate_gain(['tanh', 'relu'][use_ReLU])

        # 初始化函数
        def init_(m):
            return init(m, init_method, lambda x: nn.init.constant_(x, 0), gain=gain)

        # 获取输入通道数和尺寸
        input_channel = obs_shape[0]
        input_width = obs_shape[1]
        input_height = obs_shape[2]

        # 构建CNN序列
        self.cnn = nn.Sequential(
            # 第一个卷积层
            init_(nn.Conv2d(in_channels=input_channel,
                            out_channels=hidden_size // 2,
                            kernel_size=kernel_size,
                            stride=stride)
                  ),
            active_func,
            # 展平层
            Flatten(),
            # 第一个全连接层
            init_(nn.Linear(hidden_size // 2 * (input_width - kernel_size + stride) * (input_height - kernel_size + stride),
                            hidden_size)
                  ),
            active_func,
            # 第二个全连接层
            init_(nn.Linear(hidden_size, hidden_size)), active_func)

    def forward(self, x):
        """
        前向传播
        
        参数:
            x: 输入张量
            
        返回:
            处理后的特征
        """
        # 归一化输入
        x = x / 255.0
        # 通过CNN网络
        x = self.cnn(x)
        return x


class CNNBase(nn.Module):
    """
    CNN基础网络，用作特征提取器
    """
    def __init__(self, args, obs_shape):
        """
        初始化CNN基础网络
        
        参数:
            args: 参数对象
            obs_shape: 观察空间的形状
        """
        super(CNNBase, self).__init__()

        # 从参数中获取配置
        self._use_orthogonal = args.use_orthogonal  # 是否使用正交初始化
        self._use_ReLU = args.use_ReLU  # 是否使用ReLU激活函数
        self.hidden_size = args.hidden_size  # 隐藏层大小

        # 创建CNN层
        self.cnn = CNNLayer(obs_shape, self.hidden_size, self._use_orthogonal, self._use_ReLU)

    def forward(self, x):
        """
        前向传播
        
        参数:
            x: 输入张量
            
        返回:
            处理后的特征
        """
        x = self.cnn(x)
        return x
